import numpy as np
import random, math, os, copy
import pandas as pd

COUNT_MAX = 520

class PSO(object):
    def __init__(self, fixed_points, end, dists, costs, start=0, no_points=200, n_particles=1000, k=20, iter=1000):
        """
        :param n_particles: 种群大小
        :param k: 多种群搜索top k条路径
        :param iter: 迭代次数
        :param c: 学习参数, 范围(0, 2), 一般设置为2
        :param phi: 学习参数, 范围(0, 1)
        """
        self.n_particles = n_particles  # 种群大小
        self.k = k  # top k条路径
        self.iter = iter  # 迭代次数
        self.w = random.random()  # 惯性权重
        self.c_1 = self.c_2 = 2  # 粒子学习参数
        self.phi_1 = random.random()  # 粒子学习参数
        self.phi_2 = random.random()  # 粒子学习参数
        self.phi_3 = random.random()  # 粒子学习参数

        self.no_points = no_points  # 数据点的个数
        self.start = start
        self.fixed_points = fixed_points
        self.end = end
        # 备选点
        self.backup = [i for i in range(no_points) if i != start and i != end and i not in fixed_points]
        # 分段的节点
        self.parts = list()
        self.parts.append(start)
        self.parts.extend(fixed_points)
        self.parts.append(end)

        # 粒子的维度
        self.dims = len(self.parts) - 1

        self.costs = costs
        self.dists = dists

    def init_swarm(self):
        """
        初始化种群， 维度为3
        :param start: 开始节点, int
        :param fixed_points: 固定节点, list of int
        :param end: 结束节点, int
        :return: k个种群
        """
        swarm = []
        for i in range(self.n_particles):
            # init the paticles
            flag = True
            part_lens = []
            while flag:
                part_lens = []
                # 初始化每个分段的点的个数，之和大于20
                for j in range(self.dims):
                    part_lens.append(random.randint(5, 15))
                if sum(part_lens) >= 20:
                    flag = False
            choosed = []
            for j in range(self.dims):
                b = random.sample([i for i in self.backup if i not in choosed], part_lens[j])
                choosed.append(b)
                tmp = []
                tmp.append(self.parts[j])
                tmp.extend(b)
                tmp.append(self.parts[j+1])
                swarm.append(tmp)

        paticles_dicts = []
        for i in range(0, len(swarm), 3):
            tmp = {}
            tmp['p'] = self.concats4parts([swarm[i], swarm[i + 1], swarm[i + 2]])
            tmp['parts'] = [swarm[i], swarm[i + 1], swarm[i + 2]]
            tmp['fitness'] = self.func(tmp['p'])
            paticles_dicts.append(tmp)
        # sorted
        paticles_dicts = sorted(paticles_dicts, key=lambda x: x['fitness'])
        # divided
        divide_ps = []
        for j in range(self.k):
            divide_ps.append([paticles_dicts[i] for i in range(j, self.n_particles, self.k)])
        return divide_ps

    def concats4parts(self, parts):
        res = []
        for i in range(len(parts)):
            if i == 0:
                res.extend(parts[i])
            else:
                res.extend([parts[i][j] for j in range(len(parts[i])) if j != 0])
        return res

    def func(self, path):
        # 传入路径以及每条弧线的成本
        # path 路径, costs 为矩阵
        res = 0
        for i in range(len(path)-1):
            for j in range(i+1, len(path)):
                res += self.costs[i][j]
        return res

    def run(self):
        particles = self.init_swarm()

        # 粒子的更新依据位置交换进行
        for t in range(self.iter):

            pbest = copy.deepcopy(particles)  # 保存上一次作为备选
            for j in range(len(particles)):  # 总群
                sub_gbest = particles[j][0]  # 默认是排序好的，所以子种群最好的为序号为1的粒子
                # print('the {}th gbest fitness of sub_swarm: {}\n'.format(j, sub_gbest['fitness']))
                # print('the {}th gbest len-{} path of sub_swarm: {}\n'.format(j, len(sub_gbest['p']), sub_gbest['p']))
                # 更新位置
                for i in range(len(particles[j])):
                    for k in range(self.dims):
                        rand = random.random()
                        if rand < self.phi_1:
                            particles[j][i]['parts'][k] = copy.deepcopy(pbest[j][i]['parts'][k])
                        elif self.phi_1 <= rand and rand < self.phi_2:
                            particles[j][i]['parts'][k] = copy.deepcopy(sub_gbest['parts'][k])
                        elif rand >= self.phi_1 and rand >= self.phi_3:
                            continue
                    particles[j][i]['p'] = copy.deepcopy(self.concats4parts(particles[j][i]['parts']))
                    particles[j][i]['fitness'] = copy.deepcopy(self.func(particles[j][i]['p']))

                # 每次迭代外将子种群排序
                particles[j] = sorted(particles[j], key=lambda x: x['fitness'])

        return particles

    def calculate_path_distance(self, path):
        sum = 0.0
        for i in range(1, len(path)-1):
            sum += self.dists[path[i-1]][path[i]]
        return sum

    def compare_paths(self, path_one, path_two):
        return self.calculate_path_distance(path_one) > self.calculate_path_distance(path_two)

    def get_reverse_path(self, path):
        start = random.randint(1, len(path) - 2)
        while True:
            end = random.randint(1, len(path) - 2)
            if np.abs(start - end) > 1:
                break

        if start > end:
            path[end: start + 1] = path[end: start + 1][::-1]
            return path
        else:
            path[start: end + 1] = path[start: end + 1][::-1]
            return path

    # 不断优化，得到一个最终的最短的路径
    def update_path(self, path):
        count = 0
        while count < COUNT_MAX:
            reverse_path = self.get_reverse_path(path.copy())
            if self.compare_paths(path, reverse_path):
                count = 0
                path = reverse_path
            else:
                count += 1
        return path

    def shortest(self):
        particles = self.run()
        print(particles)
        res = []
        for i in range(len(particles)):
            sub_gbest = particles[i][0]
            res.append(self.update_path(sub_gbest['p']))
        return res

if __name__ == '__main__':
    ROOT_PATH = os.getcwd()
    DATA_PATH = os.path.join(ROOT_PATH, 'data')

    name = '附件3.xlsx'

    data = pd.read_excel(io=os.path.join(DATA_PATH, name))
    data.pop('备注')
    # print(data)

    dists = np.zeros((len(data['X坐标(米)']), len(data['X坐标(米)'])))
    for i in range(len(data['X坐标(米)'])):
        for j in range(i + 1, len(data['X坐标(米)']), 1):
            dists[i][j] = dists[j][i] = math.sqrt(
                (data['X坐标(米)'][i] - data['X坐标(米)'][j]) ** 2 + (data['Y坐标(米)'][i] - data['Y坐标(米)'][j]) ** 2 + (
                            data['Z坐标(米）'][i] - data['Z坐标(米）'][j]) ** 2) / 1000
    e_cost = dists / 6  # 边的时间花费
    enroll_cost = data['难度(分钟)'].to_numpy() / 60  # 打卡的时间花费
    row, col = e_cost.shape
    c_a = []  # 成本
    for i in range(row):
        for j in range(col):
            if i == j:
                c_a.append(0)
            else:
                e_cost[i][j] += enroll_cost[j]
                c_a.append(e_cost[i][j])

    p = PSO([182, 194], 174, dists, e_cost)    # 147, 18
    # print(p.init_swarm())
    print(p.shortest())

